

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>torch.optim.adam &mdash; pytorch_tabnet  documentation</title>
  

  
  <link rel="stylesheet" href="../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/graphviz.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/./default.css" type="text/css" />

  
  
  
  

  
  <!--[if lt IE 9]>
    <script src="../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../" src="../../../_static/documentation_options.js"></script>
        <script src="../../../_static/jquery.js"></script>
        <script src="../../../_static/underscore.js"></script>
        <script src="../../../_static/doctools.js"></script>
        <script src="../../../_static/language_data.js"></script>
    
    <script type="text/javascript" src="../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../genindex.html" />
    <link rel="search" title="Search" href="../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../index.html" class="icon icon-home" alt="Documentation Home"> pytorch_tabnet
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p><span class="caption-text">Contents:</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html">README</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#what-is-new">What is new ?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#useful-links">Useful links</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../index.html">pytorch_tabnet</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../index.html">Module code</a> &raquo;</li>
        
      <li>torch.optim.adam</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for torch.optim.adam</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">collections</span> <span class="kn">import</span> <span class="n">defaultdict</span>
<span class="kn">import</span> <span class="nn">math</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">cast</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Dict</span><span class="p">,</span> <span class="n">Tuple</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">Tensor</span>
<span class="kn">from</span> <span class="nn">.optimizer</span> <span class="kn">import</span> <span class="n">Optimizer</span><span class="p">,</span> <span class="n">_use_grad_for_differentiable</span>

<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;Adam&#39;</span><span class="p">,</span> <span class="s1">&#39;adam&#39;</span><span class="p">]</span>


<span class="c1"># TODO(crcrpar): Move this to soemwhere (e.g. torch/optim/_utils?) else when adding another fused optimizer.</span>
<span class="c1"># NOTE(crcrpar): Almost the same as `_MultiDeviceReplicator` defined in</span>
<span class="c1"># torch/cuda/amp/grad_scaler.py except for the key being str only for torch script.</span>
<span class="k">class</span> <span class="nc">_MultiDeviceReplicator</span><span class="p">:</span>
    <span class="n">main_tensor</span><span class="p">:</span> <span class="n">Tensor</span>
    <span class="n">_per_device_tensors</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">main_tensor</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">main_tensor</span> <span class="o">=</span> <span class="n">main_tensor</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_per_device_tensors</span> <span class="o">=</span> <span class="p">{</span><span class="nb">str</span><span class="p">(</span><span class="n">main_tensor</span><span class="o">.</span><span class="n">device</span><span class="p">):</span> <span class="n">main_tensor</span><span class="p">}</span>

    <span class="k">def</span> <span class="nf">get</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">device</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">device</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">_per_device_tensors</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_per_device_tensors</span><span class="p">[</span><span class="n">device</span><span class="p">]</span>
        <span class="n">tensor</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">main_tensor</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">,</span> <span class="n">non_blocking</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_per_device_tensors</span><span class="p">[</span><span class="n">device</span><span class="p">]</span> <span class="o">=</span> <span class="n">tensor</span>
        <span class="k">return</span> <span class="n">tensor</span>


<span class="c1"># todo(crcrpar): Move this to another place when adding another fused optimizer.</span>
<span class="k">def</span> <span class="nf">_get_fp16AMP_params</span><span class="p">(</span>
    <span class="o">*</span><span class="p">,</span>
    <span class="n">optimizer</span><span class="p">:</span> <span class="n">Optimizer</span><span class="p">,</span>
    <span class="n">grad_scaler</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">amp</span><span class="o">.</span><span class="n">GradScaler</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">device</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">]:</span>
    <span class="k">if</span> <span class="n">grad_scaler</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">return</span> <span class="kc">None</span>
    <span class="n">found_inf_dict</span> <span class="o">=</span> <span class="n">grad_scaler</span><span class="o">.</span><span class="n">_check_inf_per_device</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span>
    <span class="c1"># Combines found_inf tensors from all devices. As in GradScaler.update(),</span>
    <span class="c1"># tensors are combined on the scale&#39;s device, which is an arbitrary but</span>
    <span class="c1"># reasonable choice that avoids new context creation.</span>
    <span class="n">found_infs</span> <span class="o">=</span> <span class="p">[</span><span class="n">f</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">,</span> <span class="n">non_blocking</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="k">for</span> <span class="n">f</span> <span class="ow">in</span> <span class="n">found_inf_dict</span><span class="o">.</span><span class="n">values</span><span class="p">()]</span>
    <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">found_infs</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;No inf checks were recorded in _check_inf_per_device.&quot;</span>
    <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">():</span>
        <span class="n">found_inf_combined</span> <span class="o">=</span> <span class="n">cast</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="nb">sum</span><span class="p">(</span><span class="n">found_infs</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">_MultiDeviceReplicator</span><span class="p">(</span><span class="n">found_inf_combined</span><span class="p">)</span>

<span class="k">class</span> <span class="nc">Adam</span><span class="p">(</span><span class="n">Optimizer</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Implements Adam algorithm.</span>

<span class="sd">    .. math::</span>
<span class="sd">       \begin{aligned}</span>
<span class="sd">            &amp;\rule{110mm}{0.4pt}                                                                 \\</span>
<span class="sd">            &amp;\textbf{input}      : \gamma \text{ (lr)}, \beta_1, \beta_2</span>
<span class="sd">                \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)}          \\</span>
<span class="sd">            &amp;\hspace{13mm}      \lambda \text{ (weight decay)},  \: \textit{amsgrad},</span>
<span class="sd">                \:\textit{maximize}                                                              \\</span>
<span class="sd">            &amp;\textbf{initialize} :  m_0 \leftarrow 0 \text{ ( first moment)},</span>
<span class="sd">                v_0\leftarrow 0 \text{ (second moment)},\: \widehat{v_0}^{max}\leftarrow 0\\[-1.ex]</span>
<span class="sd">            &amp;\rule{110mm}{0.4pt}                                                                 \\</span>
<span class="sd">            &amp;\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\</span>

<span class="sd">            &amp;\hspace{5mm}\textbf{if} \: \textit{maximize}:                                       \\</span>
<span class="sd">            &amp;\hspace{10mm}g_t           \leftarrow   -\nabla_{\theta} f_t (\theta_{t-1})         \\</span>
<span class="sd">            &amp;\hspace{5mm}\textbf{else}                                                           \\</span>
<span class="sd">            &amp;\hspace{10mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})          \\</span>
<span class="sd">            &amp;\hspace{5mm}\textbf{if} \: \lambda \neq 0                                           \\</span>
<span class="sd">            &amp;\hspace{10mm} g_t \leftarrow g_t + \lambda  \theta_{t-1}                            \\</span>
<span class="sd">            &amp;\hspace{5mm}m_t           \leftarrow   \beta_1 m_{t-1} + (1 - \beta_1) g_t          \\</span>
<span class="sd">            &amp;\hspace{5mm}v_t           \leftarrow   \beta_2 v_{t-1} + (1-\beta_2) g^2_t          \\</span>
<span class="sd">            &amp;\hspace{5mm}\widehat{m_t} \leftarrow   m_t/\big(1-\beta_1^t \big)                   \\</span>
<span class="sd">            &amp;\hspace{5mm}\widehat{v_t} \leftarrow   v_t/\big(1-\beta_2^t \big)                   \\</span>
<span class="sd">            &amp;\hspace{5mm}\textbf{if} \: amsgrad                                                  \\</span>
<span class="sd">            &amp;\hspace{10mm}\widehat{v_t}^{max} \leftarrow \mathrm{max}(\widehat{v_t}^{max},</span>
<span class="sd">                \widehat{v_t})                                                                   \\</span>
<span class="sd">            &amp;\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/</span>
<span class="sd">                \big(\sqrt{\widehat{v_t}^{max}} + \epsilon \big)                                 \\</span>
<span class="sd">            &amp;\hspace{5mm}\textbf{else}                                                           \\</span>
<span class="sd">            &amp;\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/</span>
<span class="sd">                \big(\sqrt{\widehat{v_t}} + \epsilon \big)                                       \\</span>
<span class="sd">            &amp;\rule{110mm}{0.4pt}                                                          \\[-1.ex]</span>
<span class="sd">            &amp;\bf{return} \:  \theta_t                                                     \\[-1.ex]</span>
<span class="sd">            &amp;\rule{110mm}{0.4pt}                                                          \\[-1.ex]</span>
<span class="sd">       \end{aligned}</span>

<span class="sd">    For further details regarding the algorithm we refer to `Adam: A Method for Stochastic Optimization`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        params (iterable): iterable of parameters to optimize or dicts defining</span>
<span class="sd">            parameter groups</span>
<span class="sd">        lr (float, optional): learning rate (default: 1e-3)</span>
<span class="sd">        betas (Tuple[float, float], optional): coefficients used for computing</span>
<span class="sd">            running averages of gradient and its square (default: (0.9, 0.999))</span>
<span class="sd">        eps (float, optional): term added to the denominator to improve</span>
<span class="sd">            numerical stability (default: 1e-8)</span>
<span class="sd">        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)</span>
<span class="sd">        amsgrad (bool, optional): whether to use the AMSGrad variant of this</span>
<span class="sd">            algorithm from the paper `On the Convergence of Adam and Beyond`_</span>
<span class="sd">            (default: False)</span>
<span class="sd">        foreach (bool, optional): whether foreach implementation of optimizer</span>
<span class="sd">            is used (default: None)</span>
<span class="sd">        maximize (bool, optional): maximize the params based on the objective, instead of</span>
<span class="sd">            minimizing (default: False)</span>
<span class="sd">        capturable (bool, optional): whether this instance is safe to capture in a CUDA graph.</span>
<span class="sd">            Passing True can impair ungraphed performance, so if you don&#39;t intend to</span>
<span class="sd">            graph capture this instance, leave it False (default: False)</span>
<span class="sd">        fused (bool, optional): whether fused implementation of optimizer is used.</span>
<span class="sd">            Currently, `torch.float64`, `torch.float32`, `torch.float16`, and `torch.bfloat16`</span>
<span class="sd">            are supported. (default: False)</span>

<span class="sd">    .. _Adam\: A Method for Stochastic Optimization:</span>
<span class="sd">        https://arxiv.org/abs/1412.6980</span>
<span class="sd">    .. _On the Convergence of Adam and Beyond:</span>
<span class="sd">        https://openreview.net/forum?id=ryQu7f-RZ</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">lr</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">,</span> <span class="n">betas</span><span class="o">=</span><span class="p">(</span><span class="mf">0.9</span><span class="p">,</span> <span class="mf">0.999</span><span class="p">),</span> <span class="n">eps</span><span class="o">=</span><span class="mf">1e-8</span><span class="p">,</span>
                 <span class="n">weight_decay</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">amsgrad</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="o">*</span><span class="p">,</span> <span class="n">foreach</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                 <span class="n">maximize</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">capturable</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                 <span class="n">differentiable</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">fused</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">lr</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid learning rate: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">lr</span><span class="p">))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">eps</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid epsilon value: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">eps</span><span class="p">))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">betas</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&lt;</span> <span class="mf">1.0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid beta parameter at index 0: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">betas</span><span class="p">[</span><span class="mi">0</span><span class="p">]))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">betas</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">&lt;</span> <span class="mf">1.0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid beta parameter at index 1: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">betas</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">weight_decay</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid weight_decay value: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">weight_decay</span><span class="p">))</span>
        <span class="n">defaults</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">lr</span><span class="o">=</span><span class="n">lr</span><span class="p">,</span> <span class="n">betas</span><span class="o">=</span><span class="n">betas</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span>
                        <span class="n">weight_decay</span><span class="o">=</span><span class="n">weight_decay</span><span class="p">,</span> <span class="n">amsgrad</span><span class="o">=</span><span class="n">amsgrad</span><span class="p">,</span>
                        <span class="n">maximize</span><span class="o">=</span><span class="n">maximize</span><span class="p">,</span> <span class="n">foreach</span><span class="o">=</span><span class="n">foreach</span><span class="p">,</span> <span class="n">capturable</span><span class="o">=</span><span class="n">capturable</span><span class="p">,</span>
                        <span class="n">differentiable</span><span class="o">=</span><span class="n">differentiable</span><span class="p">,</span> <span class="n">fused</span><span class="o">=</span><span class="n">fused</span><span class="p">)</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">Adam</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="n">defaults</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">fused</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">differentiable</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;`fused` cannot be `differentiable`&quot;</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_step_supports_amp_scaling</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="c1"># TODO(crcrpar): [low prec params &amp; their higher prec copy]</span>
            <span class="c1"># Suppor AMP with FP16/BF16 model params which would need</span>
            <span class="c1"># higher prec copy of params to do update math in higher prec to</span>
            <span class="c1"># alleviate the loss of information.</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">all</span><span class="p">(</span>
                <span class="n">p</span><span class="o">.</span><span class="n">is_cuda</span> <span class="ow">and</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_floating_point</span><span class="p">(</span><span class="n">p</span><span class="p">)</span>
                <span class="k">for</span> <span class="n">pg</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">param_groups</span> <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">pg</span><span class="p">[</span><span class="s1">&#39;params&#39;</span><span class="p">]</span>
            <span class="p">):</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;FusedAdam requires all the params to be CUDA, floating point&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">__setstate__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">state</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">__setstate__</span><span class="p">(</span><span class="n">state</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">group</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">param_groups</span><span class="p">:</span>
            <span class="n">group</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;amsgrad&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
            <span class="n">group</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;maximize&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
            <span class="n">group</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;foreach&#39;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
            <span class="n">group</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;capturable&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
            <span class="n">group</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;differentiable&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
            <span class="n">group</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;fused&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
        <span class="n">state_values</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">state</span><span class="o">.</span><span class="n">values</span><span class="p">())</span>
        <span class="n">step_is_tensor</span> <span class="o">=</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">state_values</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">)</span> <span class="ow">and</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_tensor</span><span class="p">(</span><span class="n">state_values</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="s1">&#39;step&#39;</span><span class="p">])</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">step_is_tensor</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">state_values</span><span class="p">:</span>
                <span class="n">s</span><span class="p">[</span><span class="s1">&#39;step&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="n">s</span><span class="p">[</span><span class="s1">&#39;step&#39;</span><span class="p">]))</span>

    <span class="nd">@_use_grad_for_differentiable</span>
    <span class="k">def</span> <span class="nf">step</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">closure</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="o">*</span><span class="p">,</span> <span class="n">grad_scaler</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Performs a single optimization step.</span>

<span class="sd">        Args:</span>
<span class="sd">            closure (Callable, optional): A closure that reevaluates the model</span>
<span class="sd">                and returns the loss.</span>
<span class="sd">            grad_scaler (:class:`torch.cuda.amp.GradScaler`, optional): A GradScaler which is</span>
<span class="sd">                supplied from ``grad_scaler.step(optimizer)``.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_cuda_graph_capture_health_check</span><span class="p">()</span>

        <span class="n">loss</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="n">closure</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">enable_grad</span><span class="p">():</span>
                <span class="n">loss</span> <span class="o">=</span> <span class="n">closure</span><span class="p">()</span>

        <span class="k">for</span> <span class="n">group</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">param_groups</span><span class="p">:</span>
            <span class="n">params_with_grad</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">grads</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">exp_avgs</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">exp_avg_sqs</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">max_exp_avg_sqs</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">state_steps</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">beta1</span><span class="p">,</span> <span class="n">beta2</span> <span class="o">=</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;betas&#39;</span><span class="p">]</span>

            <span class="n">grad_scale</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">found_inf</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;fused&#39;</span><span class="p">]</span> <span class="ow">and</span> <span class="n">grad_scaler</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">grad_scale</span> <span class="o">=</span> <span class="n">grad_scaler</span><span class="o">.</span><span class="n">_get_scale_async</span><span class="p">()</span>
                <span class="n">device</span> <span class="o">=</span> <span class="n">grad_scale</span><span class="o">.</span><span class="n">device</span>
                <span class="n">grad_scale</span> <span class="o">=</span> <span class="n">_MultiDeviceReplicator</span><span class="p">(</span><span class="n">grad_scale</span><span class="p">)</span>
                <span class="n">found_inf</span> <span class="o">=</span> <span class="n">_get_fp16AMP_params</span><span class="p">(</span><span class="n">optimizer</span><span class="o">=</span><span class="bp">self</span><span class="p">,</span> <span class="n">grad_scaler</span><span class="o">=</span><span class="n">grad_scaler</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">)</span>

            <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;params&#39;</span><span class="p">]:</span>
                <span class="k">if</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">params_with_grad</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">p</span><span class="p">)</span>
                    <span class="k">if</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span><span class="o">.</span><span class="n">is_sparse</span><span class="p">:</span>
                        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s1">&#39;Adam does not support sparse gradients, please consider SparseAdam instead&#39;</span><span class="p">)</span>
                    <span class="n">grads</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">p</span><span class="o">.</span><span class="n">grad</span><span class="p">)</span>

                    <span class="n">state</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">state</span><span class="p">[</span><span class="n">p</span><span class="p">]</span>
                    <span class="c1"># Lazy state initialization</span>
                    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">state</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                        <span class="n">state</span><span class="p">[</span><span class="s1">&#39;step&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span>
                            <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="mi">1</span><span class="p">,),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">float</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">p</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
                            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">defaults</span><span class="p">[</span><span class="s1">&#39;capturable&#39;</span><span class="p">]</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">defaults</span><span class="p">[</span><span class="s1">&#39;fused&#39;</span><span class="p">]</span>
                            <span class="k">else</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="mf">0.</span><span class="p">)</span>
                        <span class="p">)</span>
                        <span class="c1"># Exponential moving average of gradient values</span>
                        <span class="n">state</span><span class="p">[</span><span class="s1">&#39;exp_avg&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="n">memory_format</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">preserve_format</span><span class="p">)</span>
                        <span class="c1"># Exponential moving average of squared gradient values</span>
                        <span class="n">state</span><span class="p">[</span><span class="s1">&#39;exp_avg_sq&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="n">memory_format</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">preserve_format</span><span class="p">)</span>
                        <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;amsgrad&#39;</span><span class="p">]:</span>
                            <span class="c1"># Maintains max of all exp. moving avg. of sq. grad. values</span>
                            <span class="n">state</span><span class="p">[</span><span class="s1">&#39;max_exp_avg_sq&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="n">memory_format</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">preserve_format</span><span class="p">)</span>

                    <span class="n">exp_avgs</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="s1">&#39;exp_avg&#39;</span><span class="p">])</span>
                    <span class="n">exp_avg_sqs</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="s1">&#39;exp_avg_sq&#39;</span><span class="p">])</span>

                    <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;amsgrad&#39;</span><span class="p">]:</span>
                        <span class="n">max_exp_avg_sqs</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="s1">&#39;max_exp_avg_sq&#39;</span><span class="p">])</span>
                    <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;differentiable&#39;</span><span class="p">]</span> <span class="ow">and</span> <span class="n">state</span><span class="p">[</span><span class="s1">&#39;step&#39;</span><span class="p">]</span><span class="o">.</span><span class="n">requires_grad</span><span class="p">:</span>
                        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s1">&#39;`requires_grad` is not supported for `step` in differentiable mode&#39;</span><span class="p">)</span>
                    <span class="n">state_steps</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="s1">&#39;step&#39;</span><span class="p">])</span>

            <span class="n">adam</span><span class="p">(</span><span class="n">params_with_grad</span><span class="p">,</span>
                 <span class="n">grads</span><span class="p">,</span>
                 <span class="n">exp_avgs</span><span class="p">,</span>
                 <span class="n">exp_avg_sqs</span><span class="p">,</span>
                 <span class="n">max_exp_avg_sqs</span><span class="p">,</span>
                 <span class="n">state_steps</span><span class="p">,</span>
                 <span class="n">amsgrad</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;amsgrad&#39;</span><span class="p">],</span>
                 <span class="n">beta1</span><span class="o">=</span><span class="n">beta1</span><span class="p">,</span>
                 <span class="n">beta2</span><span class="o">=</span><span class="n">beta2</span><span class="p">,</span>
                 <span class="n">lr</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;lr&#39;</span><span class="p">],</span>
                 <span class="n">weight_decay</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;weight_decay&#39;</span><span class="p">],</span>
                 <span class="n">eps</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;eps&#39;</span><span class="p">],</span>
                 <span class="n">maximize</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;maximize&#39;</span><span class="p">],</span>
                 <span class="n">foreach</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;foreach&#39;</span><span class="p">],</span>
                 <span class="n">capturable</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;capturable&#39;</span><span class="p">],</span>
                 <span class="n">differentiable</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;differentiable&#39;</span><span class="p">],</span>
                 <span class="n">fused</span><span class="o">=</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;fused&#39;</span><span class="p">],</span>
                 <span class="n">grad_scale</span><span class="o">=</span><span class="n">grad_scale</span><span class="p">,</span>
                 <span class="n">found_inf</span><span class="o">=</span><span class="n">found_inf</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">loss</span>


<span class="k">def</span> <span class="nf">adam</span><span class="p">(</span><span class="n">params</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
         <span class="n">grads</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
         <span class="n">exp_avgs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
         <span class="n">exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
         <span class="n">max_exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
         <span class="n">state_steps</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
         <span class="c1"># kwonly args with defaults are not supported by functions compiled with torchscript issue #70627</span>
         <span class="c1"># setting this as kwarg for now as functional API is compiled by torch/distributed/optim</span>
         <span class="n">foreach</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
         <span class="n">capturable</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
         <span class="n">differentiable</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
         <span class="n">fused</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
         <span class="n">grad_scale</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
         <span class="n">found_inf</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
         <span class="o">*</span><span class="p">,</span>
         <span class="n">amsgrad</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
         <span class="n">beta1</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
         <span class="n">beta2</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
         <span class="n">lr</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
         <span class="n">weight_decay</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
         <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
         <span class="n">maximize</span><span class="p">:</span> <span class="nb">bool</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Functional API that performs Adam algorithm computation.</span>
<span class="sd">    See :class:`~torch.optim.Adam` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="nb">all</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">state_steps</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;API has changed, `state_steps` argument must contain a list of singleton tensors&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">foreach</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># Placeholder for more complex foreach logic to be added when value is not set</span>
        <span class="n">foreach</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">foreach</span> <span class="ow">and</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">is_scripting</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s1">&#39;torch.jit.script not supported with foreach optimizers&#39;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">foreach</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">is_scripting</span><span class="p">():</span>
        <span class="n">func</span> <span class="o">=</span> <span class="n">_multi_tensor_adam</span>
    <span class="k">elif</span> <span class="n">fused</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">is_scripting</span><span class="p">():</span>
        <span class="n">func</span> <span class="o">=</span> <span class="n">_fused_adam</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">func</span> <span class="o">=</span> <span class="n">_single_tensor_adam</span>

    <span class="n">func</span><span class="p">(</span><span class="n">params</span><span class="p">,</span>
         <span class="n">grads</span><span class="p">,</span>
         <span class="n">exp_avgs</span><span class="p">,</span>
         <span class="n">exp_avg_sqs</span><span class="p">,</span>
         <span class="n">max_exp_avg_sqs</span><span class="p">,</span>
         <span class="n">state_steps</span><span class="p">,</span>
         <span class="n">amsgrad</span><span class="o">=</span><span class="n">amsgrad</span><span class="p">,</span>
         <span class="n">beta1</span><span class="o">=</span><span class="n">beta1</span><span class="p">,</span>
         <span class="n">beta2</span><span class="o">=</span><span class="n">beta2</span><span class="p">,</span>
         <span class="n">lr</span><span class="o">=</span><span class="n">lr</span><span class="p">,</span>
         <span class="n">weight_decay</span><span class="o">=</span><span class="n">weight_decay</span><span class="p">,</span>
         <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span>
         <span class="n">maximize</span><span class="o">=</span><span class="n">maximize</span><span class="p">,</span>
         <span class="n">capturable</span><span class="o">=</span><span class="n">capturable</span><span class="p">,</span>
         <span class="n">differentiable</span><span class="o">=</span><span class="n">differentiable</span><span class="p">,</span>
         <span class="n">grad_scale</span><span class="o">=</span><span class="n">grad_scale</span><span class="p">,</span>
         <span class="n">found_inf</span><span class="o">=</span><span class="n">found_inf</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_single_tensor_adam</span><span class="p">(</span><span class="n">params</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                        <span class="n">grads</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                        <span class="n">exp_avgs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                        <span class="n">exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                        <span class="n">max_exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                        <span class="n">state_steps</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                        <span class="n">grad_scale</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">],</span>
                        <span class="n">found_inf</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">],</span>
                        <span class="o">*</span><span class="p">,</span>
                        <span class="n">amsgrad</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
                        <span class="n">beta1</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                        <span class="n">beta2</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                        <span class="n">lr</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                        <span class="n">weight_decay</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                        <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                        <span class="n">maximize</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
                        <span class="n">capturable</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
                        <span class="n">differentiable</span><span class="p">:</span> <span class="nb">bool</span><span class="p">):</span>

    <span class="k">assert</span> <span class="n">grad_scale</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">found_inf</span> <span class="ow">is</span> <span class="kc">None</span>

    <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">param</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">params</span><span class="p">):</span>

        <span class="n">grad</span> <span class="o">=</span> <span class="n">grads</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">maximize</span> <span class="k">else</span> <span class="o">-</span><span class="n">grads</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
        <span class="n">exp_avg</span> <span class="o">=</span> <span class="n">exp_avgs</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
        <span class="n">exp_avg_sq</span> <span class="o">=</span> <span class="n">exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
        <span class="n">step_t</span> <span class="o">=</span> <span class="n">state_steps</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>

        <span class="k">if</span> <span class="n">capturable</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">param</span><span class="o">.</span><span class="n">is_cuda</span> <span class="ow">and</span> <span class="n">step_t</span><span class="o">.</span><span class="n">is_cuda</span><span class="p">,</span> <span class="s2">&quot;If capturable=True, params and state_steps must be CUDA tensors.&quot;</span>

        <span class="c1"># update step</span>
        <span class="n">step_t</span> <span class="o">+=</span> <span class="mi">1</span>

        <span class="k">if</span> <span class="n">weight_decay</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">grad</span> <span class="o">=</span> <span class="n">grad</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">param</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="n">weight_decay</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_complex</span><span class="p">(</span><span class="n">param</span><span class="p">):</span>
            <span class="n">grad</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">grad</span><span class="p">)</span>
            <span class="n">exp_avg</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">exp_avg</span><span class="p">)</span>
            <span class="n">exp_avg_sq</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">exp_avg_sq</span><span class="p">)</span>
            <span class="n">param</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">param</span><span class="p">)</span>

        <span class="c1"># Decay the first and second moment running average coefficient</span>
        <span class="n">exp_avg</span><span class="o">.</span><span class="n">mul_</span><span class="p">(</span><span class="n">beta1</span><span class="p">)</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="n">grad</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mi">1</span> <span class="o">-</span> <span class="n">beta1</span><span class="p">)</span>
        <span class="n">exp_avg_sq</span><span class="o">.</span><span class="n">mul_</span><span class="p">(</span><span class="n">beta2</span><span class="p">)</span><span class="o">.</span><span class="n">addcmul_</span><span class="p">(</span><span class="n">grad</span><span class="p">,</span> <span class="n">grad</span><span class="o">.</span><span class="n">conj</span><span class="p">(),</span> <span class="n">value</span><span class="o">=</span><span class="mi">1</span> <span class="o">-</span> <span class="n">beta2</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">capturable</span> <span class="ow">or</span> <span class="n">differentiable</span><span class="p">:</span>
            <span class="n">step</span> <span class="o">=</span> <span class="n">step_t</span>

            <span class="c1"># 1 - beta1 ** step can&#39;t be captured in a CUDA graph, even if step is a CUDA tensor</span>
            <span class="c1"># (incurs &quot;RuntimeError: CUDA error: operation not permitted when stream is capturing&quot;)</span>
            <span class="n">bias_correction1</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">beta1</span><span class="p">,</span> <span class="n">step</span><span class="p">)</span>
            <span class="n">bias_correction2</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">beta2</span><span class="p">,</span> <span class="n">step</span><span class="p">)</span>

            <span class="n">step_size</span> <span class="o">=</span> <span class="n">lr</span> <span class="o">/</span> <span class="n">bias_correction1</span>
            <span class="n">step_size_neg</span> <span class="o">=</span> <span class="n">step_size</span><span class="o">.</span><span class="n">neg</span><span class="p">()</span>

            <span class="n">bias_correction2_sqrt</span> <span class="o">=</span> <span class="n">bias_correction2</span><span class="o">.</span><span class="n">sqrt</span><span class="p">()</span>

            <span class="k">if</span> <span class="n">amsgrad</span><span class="p">:</span>
                <span class="c1"># Maintains the maximum of all 2nd moment running avg. till now</span>
                <span class="k">if</span> <span class="n">differentiable</span><span class="p">:</span>
                    <span class="n">max_exp_avg_sqs_i</span> <span class="o">=</span> <span class="n">max_exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">max_exp_avg_sqs_i</span> <span class="o">=</span> <span class="n">max_exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
                <span class="n">max_exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">copy_</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">maximum</span><span class="p">(</span><span class="n">max_exp_avg_sqs_i</span><span class="p">,</span> <span class="n">exp_avg_sq</span><span class="p">))</span>
                <span class="c1"># Uses the max. for normalizing running avg. of gradient</span>
                <span class="c1"># Folds in (admittedly ugly) 1-elem step_size math here to avoid extra param-set-sized read+write</span>
                <span class="c1"># (can&#39;t fold it into addcdiv_ below because addcdiv_ requires value is a Number, not a Tensor)</span>
                <span class="n">denom</span> <span class="o">=</span> <span class="p">(</span><span class="n">max_exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">sqrt</span><span class="p">()</span> <span class="o">/</span> <span class="p">(</span><span class="n">bias_correction2_sqrt</span> <span class="o">*</span> <span class="n">step_size_neg</span><span class="p">))</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="n">eps</span> <span class="o">/</span> <span class="n">step_size_neg</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">denom</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp_avg_sq</span><span class="o">.</span><span class="n">sqrt</span><span class="p">()</span> <span class="o">/</span> <span class="p">(</span><span class="n">bias_correction2_sqrt</span> <span class="o">*</span> <span class="n">step_size_neg</span><span class="p">))</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="n">eps</span> <span class="o">/</span> <span class="n">step_size_neg</span><span class="p">)</span>

            <span class="n">param</span><span class="o">.</span><span class="n">addcdiv_</span><span class="p">(</span><span class="n">exp_avg</span><span class="p">,</span> <span class="n">denom</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">step</span> <span class="o">=</span> <span class="n">step_t</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>

            <span class="n">bias_correction1</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">beta1</span> <span class="o">**</span> <span class="n">step</span>
            <span class="n">bias_correction2</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">beta2</span> <span class="o">**</span> <span class="n">step</span>

            <span class="n">step_size</span> <span class="o">=</span> <span class="n">lr</span> <span class="o">/</span> <span class="n">bias_correction1</span>

            <span class="n">bias_correction2_sqrt</span> <span class="o">=</span> <span class="n">math</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">bias_correction2</span><span class="p">)</span>

            <span class="k">if</span> <span class="n">amsgrad</span><span class="p">:</span>
                <span class="c1"># Maintains the maximum of all 2nd moment running avg. till now</span>
                <span class="n">torch</span><span class="o">.</span><span class="n">maximum</span><span class="p">(</span><span class="n">max_exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="n">exp_avg_sq</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">max_exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
                <span class="c1"># Use the max. for normalizing running avg. of gradient</span>
                <span class="n">denom</span> <span class="o">=</span> <span class="p">(</span><span class="n">max_exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">sqrt</span><span class="p">()</span> <span class="o">/</span> <span class="n">bias_correction2_sqrt</span><span class="p">)</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="n">eps</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">denom</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp_avg_sq</span><span class="o">.</span><span class="n">sqrt</span><span class="p">()</span> <span class="o">/</span> <span class="n">bias_correction2_sqrt</span><span class="p">)</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="n">eps</span><span class="p">)</span>

            <span class="n">param</span><span class="o">.</span><span class="n">addcdiv_</span><span class="p">(</span><span class="n">exp_avg</span><span class="p">,</span> <span class="n">denom</span><span class="p">,</span> <span class="n">value</span><span class="o">=-</span><span class="n">step_size</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_multi_tensor_adam</span><span class="p">(</span><span class="n">params</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                       <span class="n">grads</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                       <span class="n">exp_avgs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                       <span class="n">exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                       <span class="n">max_exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                       <span class="n">state_steps</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
                       <span class="n">grad_scale</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">],</span>
                       <span class="n">found_inf</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">],</span>
                       <span class="o">*</span><span class="p">,</span>
                       <span class="n">amsgrad</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
                       <span class="n">beta1</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                       <span class="n">beta2</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                       <span class="n">lr</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                       <span class="n">weight_decay</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                       <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                       <span class="n">maximize</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
                       <span class="n">capturable</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
                       <span class="n">differentiable</span><span class="p">:</span> <span class="nb">bool</span><span class="p">):</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">params</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="n">capturable</span><span class="p">:</span>
        <span class="k">assert</span> <span class="nb">all</span><span class="p">(</span><span class="n">p</span><span class="o">.</span><span class="n">is_cuda</span> <span class="ow">and</span> <span class="n">step</span><span class="o">.</span><span class="n">is_cuda</span> <span class="k">for</span> <span class="n">p</span><span class="p">,</span> <span class="n">step</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="n">state_steps</span><span class="p">)),</span> \
            <span class="s2">&quot;If capturable=True, params and state_steps must be CUDA tensors.&quot;</span>

    <span class="k">assert</span> <span class="n">grad_scale</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">found_inf</span> <span class="ow">is</span> <span class="kc">None</span>

    <span class="k">if</span> <span class="n">maximize</span><span class="p">:</span>
        <span class="n">grads</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_neg</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">grads</span><span class="p">))</span>  <span class="c1"># type: ignore[assignment]</span>

    <span class="k">assert</span> <span class="ow">not</span> <span class="n">differentiable</span><span class="p">,</span> <span class="s2">&quot;_foreach ops don&#39;t support autograd&quot;</span>
    <span class="c1"># Handle complex parameters</span>
    <span class="n">grads</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_complex</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">else</span> <span class="n">x</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">grads</span><span class="p">]</span>
    <span class="n">exp_avgs</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_complex</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">else</span> <span class="n">x</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">exp_avgs</span><span class="p">]</span>
    <span class="n">exp_avg_sqs</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_complex</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">else</span> <span class="n">x</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">exp_avg_sqs</span><span class="p">]</span>
    <span class="n">params_</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">view_as_real</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_complex</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">else</span> <span class="n">x</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">params</span><span class="p">]</span>

    <span class="c1"># update steps</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_add_</span><span class="p">(</span><span class="n">state_steps</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">weight_decay</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_add_</span><span class="p">(</span><span class="n">grads</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="n">weight_decay</span><span class="p">)</span>

    <span class="c1"># Decay the first and second moment running average coefficient</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_mul_</span><span class="p">(</span><span class="n">exp_avgs</span><span class="p">,</span> <span class="n">beta1</span><span class="p">)</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_add_</span><span class="p">(</span><span class="n">exp_avgs</span><span class="p">,</span> <span class="n">grads</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mi">1</span> <span class="o">-</span> <span class="n">beta1</span><span class="p">)</span>

    <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_mul_</span><span class="p">(</span><span class="n">exp_avg_sqs</span><span class="p">,</span> <span class="n">beta2</span><span class="p">)</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_addcmul_</span><span class="p">(</span><span class="n">exp_avg_sqs</span><span class="p">,</span> <span class="n">grads</span><span class="p">,</span> <span class="n">grads</span><span class="p">,</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">beta2</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">capturable</span><span class="p">:</span>
        <span class="c1"># TODO: use foreach_pow if/when foreach_pow is added</span>
        <span class="n">bias_correction1</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">beta1</span><span class="p">,</span> <span class="n">step</span><span class="p">)</span> <span class="k">for</span> <span class="n">step</span> <span class="ow">in</span> <span class="n">state_steps</span><span class="p">]</span>
        <span class="n">bias_correction2</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">beta2</span><span class="p">,</span> <span class="n">step</span><span class="p">)</span> <span class="k">for</span> <span class="n">step</span> <span class="ow">in</span> <span class="n">state_steps</span><span class="p">]</span>
        <span class="c1"># foreach_sub doesn&#39;t allow a scalar as the first arg</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_sub_</span><span class="p">(</span><span class="n">bias_correction1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_sub_</span><span class="p">(</span><span class="n">bias_correction2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_neg_</span><span class="p">(</span><span class="n">bias_correction1</span><span class="p">)</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_neg_</span><span class="p">(</span><span class="n">bias_correction2</span><span class="p">)</span>

        <span class="c1"># foreach_div doesn&#39;t allow a scalar as the first arg</span>
        <span class="n">step_size</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_div</span><span class="p">(</span><span class="n">bias_correction1</span><span class="p">,</span> <span class="n">lr</span><span class="p">)</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_reciprocal_</span><span class="p">(</span><span class="n">step_size</span><span class="p">)</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_neg_</span><span class="p">(</span><span class="n">step_size</span><span class="p">)</span>

        <span class="n">bias_correction2_sqrt</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_sqrt</span><span class="p">(</span><span class="n">bias_correction2</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">amsgrad</span><span class="p">:</span>
            <span class="c1"># Maintains the maximum of all 2nd moment running avg. till now</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_maximum_</span><span class="p">(</span><span class="n">max_exp_avg_sqs</span><span class="p">,</span> <span class="n">exp_avg_sqs</span><span class="p">)</span>  <span class="c1"># type: ignore[assignment]</span>

            <span class="c1"># Use the max. for normalizing running avg. of gradient</span>
            <span class="n">max_exp_avg_sq_sqrt</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_sqrt</span><span class="p">(</span><span class="n">max_exp_avg_sqs</span><span class="p">)</span>
            <span class="c1"># Folds in (admittedly ugly) 1-elem step_size math here to avoid extra param-set-sized read+write</span>
            <span class="c1"># (can&#39;t fold it into addcdiv_ below because addcdiv_ requires value is a Number, not a Tensor)</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_div_</span><span class="p">(</span><span class="n">max_exp_avg_sq_sqrt</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_mul</span><span class="p">(</span><span class="n">bias_correction2_sqrt</span><span class="p">,</span> <span class="n">step_size</span><span class="p">))</span>
            <span class="n">eps_over_step_size</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_div</span><span class="p">(</span><span class="n">step_size</span><span class="p">,</span> <span class="n">eps</span><span class="p">)</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_reciprocal_</span><span class="p">(</span><span class="n">eps_over_step_size</span><span class="p">)</span>
            <span class="n">denom</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_add</span><span class="p">(</span><span class="n">max_exp_avg_sq_sqrt</span><span class="p">,</span> <span class="n">eps_over_step_size</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">exp_avg_sq_sqrt</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_sqrt</span><span class="p">(</span><span class="n">exp_avg_sqs</span><span class="p">)</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_div_</span><span class="p">(</span><span class="n">exp_avg_sq_sqrt</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_mul</span><span class="p">(</span><span class="n">bias_correction2_sqrt</span><span class="p">,</span> <span class="n">step_size</span><span class="p">))</span>
            <span class="n">eps_over_step_size</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_div</span><span class="p">(</span><span class="n">step_size</span><span class="p">,</span> <span class="n">eps</span><span class="p">)</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_reciprocal_</span><span class="p">(</span><span class="n">eps_over_step_size</span><span class="p">)</span>
            <span class="n">denom</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_add</span><span class="p">(</span><span class="n">exp_avg_sq_sqrt</span><span class="p">,</span> <span class="n">eps_over_step_size</span><span class="p">)</span>

        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_addcdiv_</span><span class="p">(</span><span class="n">params_</span><span class="p">,</span> <span class="n">exp_avgs</span><span class="p">,</span> <span class="n">denom</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">bias_correction1</span> <span class="o">=</span> <span class="p">[</span><span class="mi">1</span> <span class="o">-</span> <span class="n">beta1</span> <span class="o">**</span> <span class="n">step</span><span class="o">.</span><span class="n">item</span><span class="p">()</span> <span class="k">for</span> <span class="n">step</span> <span class="ow">in</span> <span class="n">state_steps</span><span class="p">]</span>
        <span class="n">bias_correction2</span> <span class="o">=</span> <span class="p">[</span><span class="mi">1</span> <span class="o">-</span> <span class="n">beta2</span> <span class="o">**</span> <span class="n">step</span><span class="o">.</span><span class="n">item</span><span class="p">()</span> <span class="k">for</span> <span class="n">step</span> <span class="ow">in</span> <span class="n">state_steps</span><span class="p">]</span>

        <span class="n">step_size</span> <span class="o">=</span> <span class="p">[(</span><span class="n">lr</span> <span class="o">/</span> <span class="n">bc</span><span class="p">)</span> <span class="o">*</span> <span class="o">-</span><span class="mi">1</span> <span class="k">for</span> <span class="n">bc</span> <span class="ow">in</span> <span class="n">bias_correction1</span><span class="p">]</span>

        <span class="n">bias_correction2_sqrt</span> <span class="o">=</span> <span class="p">[</span><span class="n">math</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">bc</span><span class="p">)</span> <span class="k">for</span> <span class="n">bc</span> <span class="ow">in</span> <span class="n">bias_correction2</span><span class="p">]</span>

        <span class="k">if</span> <span class="n">amsgrad</span><span class="p">:</span>
            <span class="c1"># Maintains the maximum of all 2nd moment running avg. till now</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_maximum_</span><span class="p">(</span><span class="n">max_exp_avg_sqs</span><span class="p">,</span> <span class="n">exp_avg_sqs</span><span class="p">)</span>

            <span class="c1"># Use the max. for normalizing running avg. of gradient</span>
            <span class="n">max_exp_avg_sq_sqrt</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_sqrt</span><span class="p">(</span><span class="n">max_exp_avg_sqs</span><span class="p">)</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_div_</span><span class="p">(</span><span class="n">max_exp_avg_sq_sqrt</span><span class="p">,</span> <span class="n">bias_correction2_sqrt</span><span class="p">)</span>
            <span class="n">denom</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_add</span><span class="p">(</span><span class="n">max_exp_avg_sq_sqrt</span><span class="p">,</span> <span class="n">eps</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">exp_avg_sq_sqrt</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_sqrt</span><span class="p">(</span><span class="n">exp_avg_sqs</span><span class="p">)</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_div_</span><span class="p">(</span><span class="n">exp_avg_sq_sqrt</span><span class="p">,</span> <span class="n">bias_correction2_sqrt</span><span class="p">)</span>
            <span class="n">denom</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_add</span><span class="p">(</span><span class="n">exp_avg_sq_sqrt</span><span class="p">,</span> <span class="n">eps</span><span class="p">)</span>

        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_addcdiv_</span><span class="p">(</span><span class="n">params_</span><span class="p">,</span> <span class="n">exp_avgs</span><span class="p">,</span> <span class="n">denom</span><span class="p">,</span> <span class="n">step_size</span><span class="p">)</span>


<span class="c1"># TODO(crcrpar): Move this to another place when adding another fused optimizer.</span>
<span class="c1"># TODO(crcrpar): Make this generic when there&#39;s more fused optimizers.</span>
<span class="c1"># TODO(crcrpar): Think of rewriting this in C++.</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">()</span>
<span class="k">def</span> <span class="nf">_group_params_by_device_and_dtype</span><span class="p">(</span>
    <span class="n">params</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">grads</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">exp_avgs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">max_exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">state_steps</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Dict</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">dtype</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]]]:</span>
    <span class="n">per_device_and_dtype_tensors</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="p">[[]</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">6</span><span class="p">)])</span>
    <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="n">step</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="n">state_steps</span><span class="p">)):</span>
        <span class="n">key</span> <span class="o">=</span> <span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">p</span><span class="o">.</span><span class="n">device</span><span class="p">),</span> <span class="n">p</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
        <span class="n">per_device_and_dtype_tensors</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">p</span><span class="p">)</span>
        <span class="n">per_device_and_dtype_tensors</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">grads</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="n">per_device_and_dtype_tensors</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">exp_avgs</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="n">per_device_and_dtype_tensors</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">3</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="k">if</span> <span class="n">max_exp_avg_sqs</span><span class="p">:</span>
            <span class="n">per_device_and_dtype_tensors</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">4</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">max_exp_avg_sqs</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="n">per_device_and_dtype_tensors</span><span class="p">[</span><span class="n">key</span><span class="p">][</span><span class="mi">5</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">step</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">per_device_and_dtype_tensors</span>


<span class="k">def</span> <span class="nf">_fused_adam</span><span class="p">(</span>
    <span class="n">params</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">grads</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">exp_avgs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">max_exp_avg_sqs</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">state_steps</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">grad_scale</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">],</span>
    <span class="n">found_inf</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_MultiDeviceReplicator</span><span class="p">],</span>
    <span class="o">*</span><span class="p">,</span>
    <span class="n">amsgrad</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
    <span class="n">beta1</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
    <span class="n">beta2</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
    <span class="n">lr</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
    <span class="n">weight_decay</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
    <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
    <span class="n">maximize</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
    <span class="n">capturable</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>  <span class="c1"># Needed for consistency.</span>
    <span class="n">differentiable</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
    <span class="n">grouped_tensors</span> <span class="o">=</span> <span class="n">_group_params_by_device_and_dtype</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="n">grads</span><span class="p">,</span> <span class="n">exp_avgs</span><span class="p">,</span> <span class="n">exp_avg_sqs</span><span class="p">,</span> <span class="n">max_exp_avg_sqs</span><span class="p">,</span> <span class="n">state_steps</span><span class="p">)</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="p">)</span> <span class="ow">in</span> <span class="n">grouped_tensors</span><span class="p">:</span>
        <span class="p">(</span>
            <span class="n">device_params</span><span class="p">,</span>
            <span class="n">device_grads</span><span class="p">,</span>
            <span class="n">device_exp_avgs</span><span class="p">,</span>
            <span class="n">device_exp_avg_sqs</span><span class="p">,</span>
            <span class="n">device_max_exp_avg_sqs</span><span class="p">,</span>
            <span class="n">device_state_steps</span><span class="p">,</span>
        <span class="p">)</span> <span class="o">=</span> <span class="n">grouped_tensors</span><span class="p">[(</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="p">)]</span>
        <span class="k">if</span> <span class="n">grad_scale</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">found_inf</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">device_grad_scale</span> <span class="o">=</span> <span class="n">grad_scale</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
            <span class="n">device_found_inf</span> <span class="o">=</span> <span class="n">found_inf</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">device_grad_scale</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">device_found_inf</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_add_</span><span class="p">(</span><span class="n">device_state_steps</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_fused_adam_</span><span class="p">(</span>
            <span class="n">device_params</span><span class="p">,</span>
            <span class="n">device_grads</span><span class="p">,</span>
            <span class="n">device_exp_avgs</span><span class="p">,</span>
            <span class="n">device_exp_avg_sqs</span><span class="p">,</span>
            <span class="n">device_max_exp_avg_sqs</span><span class="p">,</span>
            <span class="n">device_state_steps</span><span class="p">,</span>
            <span class="n">amsgrad</span><span class="o">=</span><span class="n">amsgrad</span><span class="p">,</span>
            <span class="n">lr</span><span class="o">=</span><span class="n">lr</span><span class="p">,</span>
            <span class="n">beta1</span><span class="o">=</span><span class="n">beta1</span><span class="p">,</span>
            <span class="n">beta2</span><span class="o">=</span><span class="n">beta2</span><span class="p">,</span>
            <span class="n">weight_decay</span><span class="o">=</span><span class="n">weight_decay</span><span class="p">,</span>
            <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span>
            <span class="n">maximize</span><span class="o">=</span><span class="n">maximize</span><span class="p">,</span>
            <span class="n">grad_scale</span><span class="o">=</span><span class="n">device_grad_scale</span><span class="p">,</span>
            <span class="n">found_inf</span><span class="o">=</span><span class="n">device_found_inf</span><span class="p">,</span>
        <span class="p">)</span>
        <span class="k">if</span> <span class="n">device_found_inf</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_foreach_sub_</span><span class="p">(</span><span class="n">device_state_steps</span><span class="p">,</span> <span class="p">[</span><span class="n">device_found_inf</span><span class="p">]</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">device_state_steps</span><span class="p">))</span>
</pre></div>

           </div>
           
          </div>
          <footer>
  

  <hr/>

  <div role="contentinfo">
    <p>
        
        &copy; Copyright 2019, Dreamquark

    </p>
  </div>
    
    
    
    Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>

        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>